import numpy 
from scipy import linalg
import matplotlib.pyplot as plt

pi = 0.25 # conformation 
mu = 1 # mutation rate  
Q = numpy.mat('[-3 1 1 1; 1 -3 1 1; 1 1 -3 1; 1 1 1 -3]')/4.0 * mu
beta =  1/(-pi*Q.diagonal().sum())
Q = Q * beta # normalize the Q independent of mu rates

gene1=[2, 2, 0, 1] # GGAC 
gene2=[2, 0, 1, 1] # GACC

## log likelihood calculattion
## exp(Q * t) =  P 
def llh(a, b, Q, t):
	P = linalg.expm(Q*t) 
	sum = 0
	for i in range(len(a)):
		sum += numpy.log(pi) + numpy.log(P[ a[i], b[i]]) ## log (  P(G) * P(G=>G) * P(G) * P(G=>A) ... P(C) * P(C=>C)
	return sum

t = numpy.linspace(0,2,100) # branch lengths
y = [ llh(gene1,gene2,Q,x) for x in t] # log likelihood 

## find the maximum likelihood of branch length
plt.plot(t,y)
plt.show()


